# Copyright (c) 2024 Intel Corporation
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#      http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import weakref
from typing import Any, Optional

from torch import nn


class RemovableHookHandle:
    """
    A handle to manage removable hooks that are stored in a dictionary within a PyTorch ModuleDict.
    The handle allows safe removal of specific hooks based on their unique identifiers.

    :param storage_ref: A weak reference to the storage (an nn.ModuleDict) containing the hooks.
    :param key: The key in the storage dict corresponding to a set of hooks.
    :param id: A unique identifier for the specific hook within the set.
    """

    def __init__(self, storage: nn.ModuleDict, key: str, id: str) -> None:
        self.storage_ref = weakref.ref(storage)
        self.key = key
        self.id = id

    def remove(self) -> None:
        """
        Removes the hook with the specified id from the storage, if it exists.
        If the set of hooks under a particular key becomes empty after removal, the key is also deleted.
        Steps:
        - Checks if the storage still exists (i.e., if it's not garbage collected).
        - Verifies if the key and id are present in the storage.
        - If hook exists, hook will removed.
        - If no hooks left under the key, the key is also removed from the storage.
        """
        storage: Optional[nn.ModuleDict] = self.storage_ref()
        if storage is not None and self.key in storage and self.id in storage[self.key]:
            del storage[self.key][self.id]
            if not storage[self.key]:
                del storage[self.key]


class HookStorage(nn.Module):
    """
    A module for storing and executing hooks.

    :param pre_hooks: A instance of nn.ModuleDict for storing pre-hooks.
    :param post_hooks: A instance of nn.ModuleDict for storing post-hooks.
    """

    def __init__(self) -> None:
        """
        Initialize an empty HookStorage.
        """
        super().__init__()
        self.pre_hooks: nn.ModuleDict = nn.ModuleDict()
        self.post_hooks: nn.ModuleDict = nn.ModuleDict()

    @staticmethod
    def _generate_key(op_name: str, port_id: int) -> str:
        """
        Return key for module dict that generated by operation name and port id.

        :param op_name: The operation name the hook is associated with.
        :param port_id: The port ID the hook is associated with.
        :returns: Generate key for module dict.
        """
        return f"{op_name}__{port_id}"

    @staticmethod
    def _get_next_hook_id(hooks_dict: nn.ModuleDict) -> str:
        """
        Determines the next available hook ID by finding the highest existing hook ID and incrementing it.

        :param hooks_dict: A dictionary containing existing hooks.
        :returns: The next available hook ID as a string. Starts from '0' if no hooks exist.
        """

        if not hooks_dict:
            return "0"
        return str(max([int(k) for k in hooks_dict]) + 1)

    @classmethod
    def _insert_hook(
        cls, storage_dict: nn.ModuleDict, op_name: str, port_id: int, hook: nn.Module
    ) -> RemovableHookHandle:
        """
        Inserts a hook into the storage under the appropriate key.

        :param storage_dict: The storage of hook.
        :param op_name: The operation name the hook is associated with.
        :param port_id: The port ID the hook is associated with.
        :param hook: The hook module to be stored.
        :returns: A handle that can be used to remove the hook later.
        """

        hook_key = cls._generate_key(op_name, port_id)
        if hook_key not in storage_dict:
            storage_dict[hook_key] = nn.ModuleDict()

        hook_id = cls._get_next_hook_id(storage_dict[hook_key])
        storage_dict[hook_key][hook_id] = hook

        return RemovableHookHandle(storage_dict, hook_key, hook_id)

    def register_pre_function_hook(self, op_name: str, port_id: int, hook: nn.Module) -> RemovableHookHandle:
        """
        Registers a pre-function hook to be executed before a specific operation.

        :param op_name: The operation name the hook is associated with.
        :param port_id: The port ID the hook is associated with.
        :param hook: The pre-function hook to be stored.
        :returns: A handle for removing the registered hook.
        """
        return self._insert_hook(self.pre_hooks, op_name, port_id, hook)

    def register_post_function_hook(self, op_name: str, port_id: int, hook: nn.Module) -> RemovableHookHandle:
        """
        Registers a post-function hook to be executed before a specific operation.

        :param op_name: The operation name the hook is associated with.
        :param port_id: The port ID the hook is associated with.
        :param hook: The pre-function hook to be stored.
        :returns: A handle for removing the registered hook.
        """
        return self._insert_hook(self.post_hooks, op_name, port_id, hook)

    @classmethod
    def _execute_hooks(cls, storage_dict: nn.ModuleDict, op_name: str, port_id: int, value: Any) -> Any:
        """
        Executes all hooks of a given type (pre or post) for a specific operation and port,
        passing and potentially modifying a value.

        :param storage_dict: The storage of hook.
        :param op_name: The operation name the hooks are associated with.
        :param port_id: The port ID the hooks are associated with.
        :param value: The input value to be passed through the hooks.
        :returns: The modified value after all hooks have been applied.
        """
        hook_key = cls._generate_key(op_name, port_id)
        if hook_key not in storage_dict:
            return value
        for hook in storage_dict[hook_key].values():
            value = hook(value)
        return value

    def execute_pre_function_hooks(self, op_name: str, port_id: int, value: Any) -> Any:
        """
        Executes all pre-function hooks for a given operation and port.

        :param op_name: The operation name the hooks are associated with.
        :param port_id: The port ID the hooks are associated with.
        :param value: The input value to be passed through the pre-function hooks.
        :returns: The value after all pre-function hooks have been executed.
        """
        return self._execute_hooks(self.pre_hooks, op_name, port_id, value)

    def execute_post_function_hooks(self, op_name: str, port_id: int, value: Any) -> Any:
        """
        Executes all post-function hooks for a given operation and port.

        :param op_name: The operation name the hooks are associated with.
        :param port_id: The port ID the hooks are associated with.
        :param value: The input value to be passed through the pre-function hooks.
        :returns: The value after all post-function hooks have been executed.
        """
        return self._execute_hooks(self.post_hooks, op_name, port_id, value)
